# -*- coding: utf-8 -*-


import math

from mtl.op_exec import *
from mtl.op_tran import *
from mtl.op_utils import *
from mtl.hw_defs import VARS

from programs.opcf.mm import *


def fft1d_w(N, k):
    return complex(math.cos(2 * math.pi * k / N),
                   -math.sin(2 * math.pi * k / N))


def fft1d_w_list(N):
    # The value should be the power of 2.
    assert math.log(N, 2).is_integer()
    # Get the half of the value.
    half_of_elements = int(N / 2)
    # Create the list of w.
    ws = [0 + 0j] * half_of_elements
    # Loop for all the elements for half of the elements.
    for k in range(half_of_elements):
        ws[k] = fft1d_w(N, k)
    return ws


def fft1d_w_dict(N):
    # Construct all the levels before N.
    weight_dict = {}
    # Loop until N reaches 1.
    while N > 1:
        # Save the dictionary.
        weight_dict[N] = fft1d_w_list(N)
        # Divided by 2.
        N = int(N / 2)
    # Add the fake 1 for padding.
    weight_dict[1] = [0]
    return weight_dict


def fft1d_64_weights():
    weights = fft1d_w_dict(64)
    weight_list = []
    # Based on weights, generate list.
    base = 1
    while base <= 64:
        weight_list += weights[base]
        base *= 2
    return weight_list


def fft2d_64_deploy_weight(addr_weight):
    # First deploy weight to chip, weight should be 64 long, each ALU get 16
    # elements.
    sm_addr_weights = smalloc('fft2d', 16)
    # Copy the data to softmem.
    for ii in range(VARS.num_of_alu):
        ld_v(pid=mtl_pid_all(),
             dst=sm_addr_weights + ii * VARS.alu_offset,
             src=addr_weight + ii * 16,
             count=16)
    # Construct the weight dictionary.
    return {
        2: sm_addr_weights + 1,  # W(0, 2)
        4: sm_addr_weights + 2,  # W(0, 4) -> W(1, 4)
        8: sm_addr_weights + 4,  # W(0, 8) -> W(3. 8)
        16: sm_addr_weights + 8,  # W(0, 16) -> W(7, 16)
        32: sm_addr_weights + VARS.alu_offset,  # W(0, 32) -> W(15, 32)
        640: sm_addr_weights + VARS.alu_offset * 2,  # W(0, 64) -> W(15, 64)
        641: sm_addr_weights + VARS.alu_offset * 3  # W(16, 64) -> W(31, 64)
    }


def fft1d_core(item_range, offset, index_offset, data_index,
               sm_addr_row, sm_addr_tmp, sm_addr_weight):
    for item_index in item_range:
        id_down = item_index + offset
        id_up = id_down - index_offset
        mul(pid=mtl_pid_all(), alu=mtl_alu_all(), dst=sm_addr_tmp, offset=True,
            src0=sm_addr_row + data_index[id_down], offset0=True,
            wait0=index_offset == 1,
            src1=sm_addr_weight + item_index, offset1=False)
        sub(pid=mtl_pid_all(), alu=mtl_alu_all(),
            dst=sm_addr_row + data_index[id_down], offset=True,
            src0=sm_addr_row + data_index[id_up], offset0=True,
            src1=sm_addr_tmp, offset1=True)
        add(pid=mtl_pid_all(), alu=mtl_alu_all(),
            dst=sm_addr_row + data_index[id_up], offset=True,
            src0=sm_addr_row + data_index[id_up], offset0=True,
            src1=sm_addr_tmp, offset1=True)


def fft2d_64_fft1d(sm_addr_row, sm_addr_w, sm_addr_tmp, data_index):
    index_offset = 1
    for n in range(5):
        for offset in range(index_offset, 64, index_offset * 2):
            # Called FFT2D core.
            fft1d_core(range(index_offset), offset, index_offset, data_index,
                       sm_addr_row, sm_addr_tmp, sm_addr_w[index_offset * 2])
        # Increase to next level.
        index_offset *= 2
    # For FFT1D(64), needs to be done in half and half.
    fft1d_core(range(16), 32, 32, data_index,
               sm_addr_row, sm_addr_tmp, sm_addr_w[640])
    fft1d_core(range(16), 48, 32, data_index,
               sm_addr_row, sm_addr_tmp, sm_addr_w[641])


def fft2d_64_transpose(sm_addr_row, data_index):
    # Construct the peer-alu group (pa_group)
    pa_groups = []
    for peer_id in range(VARS.num_of_peer):
        for alu_id in range(VARS.num_of_alu):
            pa_groups.append([peer_id, alu_id])
    # Loops for all the peers.
    for row_id in range(len(pa_groups)):
        row_peer, row_alu = pa_groups[row_id]
        # Send the data to itself.
        put(dst_pid=mtl_pid(row_peer),
            dst=sm_addr_row + VARS.alu_offset * row_alu + data_index[row_id],
            src_pid=mtl_pid(row_peer),
            src=sm_addr_row + VARS.alu_offset * row_alu + data_index[row_id], src_offset=False)
        # Loop for all the rows left.
        # Target -> Current
        for ii in range(row_id + 1, 64):
            target_peer, target_alu = pa_groups[ii]
            put(dst_pid=mtl_pid(row_peer),
                dst=sm_addr_row + VARS.alu_offset * row_alu + data_index[ii],
                dst_sb=True,
                src_pid=mtl_pid(target_peer), src_alu=mtl_alu(target_alu),
                src=sm_addr_row + VARS.alu_offset * target_alu + data_index[
                    row_id], src_offset=False)
        # Current -> Target
        for ii in range(row_id + 1, 64):
            target_peer, target_alu = pa_groups[ii]
            put(dst_pid=mtl_pid(target_peer),
                dst=sm_addr_row + VARS.alu_offset * target_alu + data_index[
                    row_id], dst_sb=True,
                src_pid=mtl_pid(row_peer), src_alu=mtl_alu(row_alu),
                src=sm_addr_row + VARS.alu_offset * row_alu + data_index[ii],
                src_offset=False)
            # Remote ALU should hold and wait for data passing in.
            wait(pid=mtl_pid(target_peer), alu=mtl_alu(target_alu),
                 src=sm_addr_row + data_index[row_id] + VARS.alu_offset * target_alu,
                 src_offset=False)


def fft2d_64(addr_image, addr_weight):
    # Deploy the weight.
    sm_addr_w = fft2d_64_deploy_weight(addr_weight)
    # Deploy the image to all the chip.
    sm_addr_row = smalloc('fft2d', 64)
    addr_row = addr_image
    for peer_id in range(16):
        for alu_id in range(4):
            ld_v(pid=mtl_pid(peer_id),
                 dst=alu_id * VARS.alu_offset + sm_addr_row,
                 src=addr_row,
                 count=64)
            addr_row += 64
    # Allocate a tmp position.
    sm_addr_tmp = smalloc('fft2d', 1)
    # Get the alu data index.
    reverse_index = [0, 32, 16, 48, 8, 40, 24, 56, 4, 36, 20, 52, 12, 44, 28, 60,
                     2, 34, 18, 50, 10, 42, 26, 58, 6, 38, 22, 54, 14, 46, 30,
                     62, 1, 33, 17, 49, 9, 41, 25, 57, 5, 37, 21, 53, 13, 45, 29,
                     61, 3, 35, 19, 51, 11, 43, 27, 59, 7, 39, 23, 55, 15, 47,
                     31, 63]
    normal_index = list(range(64))
    # Do FFT1D for each row.
    fft2d_64_fft1d(sm_addr_row, sm_addr_w, sm_addr_tmp, reverse_index)
    # Transpose the matrix.
    fft2d_64_transpose(sm_addr_row, reverse_index)
    # Do FFT1D in normal index.
    fft2d_64_fft1d(sm_addr_row, sm_addr_w, sm_addr_tmp, normal_index)
    # Transfer all the data back to original position.
    put_v(dst_pid=mtl_pid_main_mem(), dst=addr_image, dst_inc=64, count=64,
          src_pid=mtl_pid_all(), src_alu=mtl_alu_all(), src=sm_addr_row,
          src_offset=True)